package nl.altenpts.net.broker;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.NoSuchAlgorithmException;
import java.security.PublicKey;
import java.security.Signature;
import java.security.SignatureException;
import java.security.spec.X509EncodedKeySpec;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;

import nl.altenpts.net.NetReactor;
/**
 * SignedMessageBroker wraps around MessageBroker and signs outgoing messages and
 * verifies incoming messages. Outgoing messages are augmented with a signature and a
 * public key which can be used to verify the signature. Incoming messages are 
 * verified using the added signature and public key.
 *  
 * @author fabian
 *
 */
public class SignedMessageBroker {
	
	/**
	 * Constructs a new SignedMessageBroker.
	 * 
	 * @param reactor the net reactor for the underlying MessageBroker
	 * @param keys KeyPair to sign outgoing messages.
	 * @param host Interface to bind to.
	 * @param address as InetSocketAddress object which has a hostname and a socket
	 * @param sigalg Signature algorithm
	 * @throws NoSuchAlgorithmException If no Signature with algorithm sigalg can be instantiated
	 *         or if KeyFactory with algorithm of public key fails
	 * @throws IOException If MessageBroker construction fails
	 * @throws InvalidKeyException If initialization of Signature with
	 * 	       PrivateKey fails
	 */
	public SignedMessageBroker(NetReactor reactor, KeyPair keys, InetSocketAddress address, String sigalg) throws NoSuchAlgorithmException, IOException, InvalidKeyException {
		
		verifier = Signature.getInstance(sigalg);
		signer   = Signature.getInstance(sigalg);
		signer.initSign(keys.getPrivate());
		keyfactory = KeyFactory.getInstance(keys.getPublic().getAlgorithm());
		
		hooks = new HashSet<SignedMessageHook>();
		
		byte[] encpubkey = keys.getPublic().getEncoded();
		pubkey = ByteBuffer.allocate(encpubkey.length + 4);
		pubkey.putInt(encpubkey.length);
		pubkey.put(encpubkey);
		pubkey.flip();
		pubkey.mark();
		
		this.msghandler = new MessageHandler() {

			public void onIncoming(IncomingMessage message) {
				ByteBuffer payload = message.getPayload();				
				byte[] encpubkey = new byte[payload.getInt()];
				payload.get(encpubkey);
				byte[] sign = new byte[payload.getInt()];
				payload.get(sign);
				
				X509EncodedKeySpec keyspec = new X509EncodedKeySpec(encpubkey);
				
				try {
					PublicKey pubkey = keyfactory.generatePublic(keyspec);
					verifier.initVerify(pubkey);					
					payload.mark();
					verifier.update(payload);
					if(!verifier.verify(sign)) {
						throw new SignatureException();						
					}
					payload.reset();
					try {
						Iterator<SignedMessageHook> hookiterator = hooks.iterator();
						while(hookiterator.hasNext()) {
							hookiterator.next().onIncoming(message, pubkey);
						}
						if(smsghandler != null) { 
							smsghandler.onIncoming(message, pubkey);
						}
					} catch(Exception ex) {
						/* hooks and handlers should never let an exception propagate
						 * to here
						 */
						assert(false);
					}
				} catch (Exception e) {					
					// signature check failed, just drop message
				}
			}
		};
		
		
		this.broker = new MessageBroker(reactor, address);
		this.broker.setHandler(this.msghandler);
		
	}
	
	/**
	 * Sends a signed message
	 * 
	 * @param message the message to send
	 */
	public void sendMessage(OutgoingMessage message) {
		byte[] sign = null;
		ByteBuffer signbuf;

		try {
			synchronized(signer) {			
				Iterator<ByteBuffer> iterator = message.getPayload().iterator();
				while(iterator.hasNext()) {
					ByteBuffer buffer = iterator.next();
					buffer.mark();

					signer.update(buffer);

					buffer.reset();
				}
				sign = signer.sign();
			}

			signbuf = ByteBuffer.allocate(sign.length + 4);
			signbuf.putInt(sign.length);
			signbuf.put(sign);
			signbuf.flip();
			message.prependContent(signbuf);
			pubkey.reset();
			message.prependContent(pubkey);

			broker.sendMessage(message);
		} catch(SignatureException e) {
			assert(false);
		}
	}
	
	/**
	 * Set a new message handler
	 * 
	 * @param handler the new handler
	 */
	public void setHandler(SignedMessageHandler handler) {
		this.smsghandler = handler;
	}
	
	/**
	 * Returns the current handler
	 * 
	 * @return current handler 
	 */
	public SignedMessageHandler getHandler() {
		return smsghandler;
	}
	
	/**
	 * Add a new hook.
	 * 
	 * @param hook The hook to add
	 */
	public void addHook(SignedMessageHook hook) {
		hooks.add(hook);
	}
	
	/**
	 * Remove hook
	 * 
	 * @param hook the hook to be removed
	 */
	public void removeHook(SignedMessageHook hook) {
		hooks.remove(hook);
	}
	
	private MessageBroker broker;
	private MessageHandler msghandler;
	private SignedMessageHandler smsghandler = null;
	private KeyFactory keyfactory;
	private Signature verifier;
	private Signature signer;
	
	private ByteBuffer pubkey;
	private Set<SignedMessageHook> hooks;
}
